import os
import numpy as np
import nibabel as nib
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import skimage.transform  # 新增导入
from scipy.ndimage import gaussian_filter, map_coordinates
import matplotlib.pyplot as plt

class MRIDataset(Dataset):
    def __init__(self, root_dir, modalities=('axi', 'sag', 'cor'), max_slices=4, crop_size=224, augment=None):
        self.root_dir = root_dir
        self.modalities = modalities
        self.max_slices = max_slices
        self.crop_size = crop_size
        self.augment = augment
        self.labels = pd.read_csv(os.path.join(root_dir, 'label.csv'),
                                  header=0,
                                  names=['p_id', 'label']).set_index('p_id')
        self.labels['label'] = self.labels['label'].astype(int) - 1  # 转换为0-based
        assert self.labels['label'].between(0, 3).all(), "标签必须在0-3范围内"
        
        self.patient_ids = []
        for pid in self.labels.index:
            if self._check_data_complete(pid):
                self.patient_ids.append(pid)

        self.class_counts = {i: 0 for i in range(4)}  # 按0-based初始化
        for pid in self.patient_ids:
            label = self.labels.loc[pid, 'label']
            self.class_counts[label] += 1

    def _check_data_complete(self, pid):
        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            if not (os.path.exists(img_path) and os.path.exists(roi_path)):
                return False
            
            try:
                img = nib.load(img_path)
                roi = nib.load(roi_path)
            except Exception as e:
                print(f"加载文件失败：{img_path} 或 {roi_path}, 错误：{e}")
                return False
            
            img_shape = img.shape[:2]
            roi_shape = roi.shape[:2]
            
            if img_shape != roi_shape:
                print(f"患者 {pid} 的 {mod} 模态图像尺寸 {img_shape} 与 ROI 尺寸 {roi_shape} 不匹配")
                return False
            
            max_allowed_ratio = 1.5
            h, w = img_shape
            if h == 0 or w == 0:
                print(f"患者 {pid} 的 {mod} 模态存在零尺寸：{img_shape}")
                return False
            current_ratio = max(h / w, w / h)
            if current_ratio > max_allowed_ratio:
                print(f"患者 {pid} 的 {mod} 模态长宽比过大：{img_shape} (比例 {current_ratio:.2f})")
                return False
            
            # 新增检查：至少存在一个包含ROI的切片
            roi_data = roi.get_fdata()
            if not np.any(roi_data > 0):
                print(f"患者 {pid} 的 {mod} 模态没有有效ROI切片")
                return False
            
        return True

    def _apply_cropping(self, image, crop_coords, is_roi=False):
        """使用给定的裁剪坐标进行裁剪、填充和调整大小"""
        y_start, y_end, x_start, x_end = crop_coords
        cropped = image[y_start:y_end, x_start:x_end]

        # 计算填充参数
        scaled_size = int(self.crop_size / 1.4)
        half_size = scaled_size // 2
        
        pad_y_before = max(0, half_size - (y_end - y_start) // 2)
        pad_y_after = max(0, half_size - (y_end - y_start + 1) // 2)
        pad_x_before = max(0, half_size - (x_end - x_start) // 2)
        pad_x_after = max(0, half_size - (x_end - x_start + 1) // 2)

        padded = np.pad(cropped,
                       ((pad_y_before, pad_y_after),
                        (pad_x_before, pad_x_after)),
                       mode='constant')

        # 选择插值方法
        order = 0 if is_roi else 1
        
        resized = skimage.transform.resize(
            padded,
            (self.crop_size, self.crop_size),
            order=order,
            preserve_range=True,
            anti_aliasing=False
        )
        return resized.astype(np.float32)

    def __len__(self):
        return len(self.patient_ids)

    def __getitem__(self, idx):
        pid = self.patient_ids[idx]
        label = self.labels.loc[pid, 'label']
        
        multimodal_img = []
        multimodal_roi = []

        for mod in self.modalities:
            img = nib.load(os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz'))
            roi = nib.load(os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz'))
            
            img_data = img.get_fdata().astype(np.float32)
            roi_data = roi.get_fdata().astype(np.float32)

            img_slices = []
            roi_slices = []
            last_crop_coords = None

            for slice_idx in range(img_data.shape[2]):
                if len(img_slices) >= self.max_slices:
                    break
                
                img_slice = img_data[..., slice_idx]
                roi_slice = roi_data[..., slice_idx]
                
                current_has_roi = np.any(roi_slice > 0)
                
                if current_has_roi or last_crop_coords is not None:
                    if current_has_roi:
                        # 计算新的裁剪坐标
                        coords = np.argwhere(roi_slice > 0)
                        min_y, min_x = coords.min(axis=0)
                        max_y, max_x = coords.max(axis=0)
                        center_y = (min_y + max_y) // 2
                        center_x = (min_x + max_x) // 2
                        scaled_size = int(self.crop_size / 1.4)
                        half_size = scaled_size // 2
                        y_start = max(0, center_y - half_size)
                        y_end = min(img_slice.shape[0], center_y + half_size)
                        x_start = max(0, center_x - half_size)
                        x_end = min(img_slice.shape[1], center_x + half_size)
                        crop_coords = (y_start, y_end, x_start, x_end)
                        last_crop_coords = crop_coords
                        
                        # 处理图像和ROI
                        img_cropped = self._apply_cropping(img_slice, crop_coords)
                        roi_cropped = self._apply_cropping(roi_slice, crop_coords, is_roi=True)
                        roi_cropped = (roi_cropped > 0.5).astype(np.float32)
                    else:
                        # 使用最后记录的坐标处理当前切片
                        img_cropped = self._apply_cropping(img_slice, last_crop_coords)
                        roi_cropped = np.zeros_like(img_cropped)
                    
                    # 归一化处理
                    slice_min = img_slice.min()
                    slice_max = img_slice.max()
                    img_cropped = (img_cropped - slice_min) / (slice_max - slice_min + 1e-8)
                    
                    # 数据增强
                    if self.augment:
                        img_cropped, roi_cropped = self.apply_augment(img_cropped, roi_cropped)
                    
                    img_slices.append(img_cropped)
                    roi_slices.append(roi_cropped)

            # 确保填充到max_slices（处理边缘情况）
            while len(img_slices) < self.max_slices:
                img_slices.append(np.zeros((self.crop_size, self.crop_size)))
                roi_slices.append(np.zeros((self.crop_size, self.crop_size)))
            
            multimodal_img.append(torch.FloatTensor(np.stack(img_slices)))
            multimodal_roi.append(torch.FloatTensor(np.stack(roi_slices)))

        return (
            torch.stack(multimodal_img),  # (B, M, D, H, W)
            torch.stack(multimodal_roi),  # (B, M, D, H, W)
            torch.tensor(label, dtype=torch.long)  # (B,)
        )

    def apply_augment(self, img, roi):
        """应用同步几何增强"""
        if np.random.rand() < 0.5:
            # 缩放增强
            scale_factor = np.random.uniform(0.9, 1.1)
            new_size = int(self.crop_size * scale_factor)
        
            img = skimage.transform.resize(
                img, (new_size, new_size),
                order=1, preserve_range=True, anti_aliasing=True
            )
            roi = skimage.transform.resize(
                roi, (new_size, new_size),
                order=0, preserve_range=True, anti_aliasing=False
            )

            # 随机裁剪/填充
            if new_size > self.crop_size:
                start = np.random.randint(0, new_size - self.crop_size)
                img = img[start:start+self.crop_size, start:start+self.crop_size]
                roi = roi[start:start+self.crop_size, start:start+self.crop_size]
            else:
                pad = self.crop_size - new_size
                img = np.pad(img, ((0, pad), (0, pad)), mode='constant')
                roi = np.pad(roi, ((0, pad), (0, pad)), mode='constant')

        # 颜色空间增强
        if np.random.rand() < 0.5:
            img = img * np.random.uniform(0.8, 1.2) + np.random.uniform(-0.1, 0.1)
            img = np.clip(img, img.min(), img.max())

            # 旋转和平移
            if np.random.rand() < 0.5:
                angle = np.random.uniform(-10, 10)
                img = skimage.transform.rotate(img, angle, order=1, preserve_range=True, mode='reflect')
                roi = skimage.transform.rotate(roi, angle, order=0, preserve_range=True, mode='reflect')

            if np.random.rand() < 0.5:
                max_offset = int(self.crop_size * 0.1)
                dy, dx = np.random.randint(-max_offset, max_offset, size=2)
                img = np.pad(img, max_offset, mode='reflect')[
                    max_offset+dy : max_offset+dy+self.crop_size,
                    max_offset+dx : max_offset+dx+self.crop_size
                ]
                roi = np.pad(roi, max_offset, mode='constant')[
                    max_offset+dy : max_offset+dy+self.crop_size,
                    max_offset+dx : max_offset+dx+self.crop_size
                ]

        # 弹性形变（新增部分）
        if np.random.rand() < 0.5:
            alpha = self.crop_size * 2  # 形变强度系数
            sigma = self.crop_size * 0.08  # 平滑系数
            random_state = np.random.RandomState(None)
            
            # 生成随机位移场
            shape = img.shape
            dx = random_state.rand(*shape) * 2 - 1  # [-1,1]区间
            dy = random_state.rand(*shape) * 2 - 1
            
            # 高斯滤波平滑位移场
            dx = gaussian_filter(dx, sigma) * alpha
            dy = gaussian_filter(dy, sigma) * alpha
            
            # 创建网格坐标
            x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
            
            # 应用位移映射
            indices = (y + dy).reshape(-1, 1), (x + dx).reshape(-1, 1)
            img = map_coordinates(img, indices, order=1, mode='reflect').reshape(shape)
            roi = map_coordinates(roi, indices, order=0, mode='constant').reshape(shape)

        return img, roi

def visualize_attention(img_batch, roi_batch, labels):
    # 获取首个样本数据
    img, roi, label = img_batch[0].numpy(), roi_batch[0].numpy(), labels[0].item()
    mod_names = ['Axial', 'Sagittal', 'Coronal']
    
    # 创建图形窗口
    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    plt.suptitle(f"Label: {label} | Slice: 0/{img.shape[1]-1}", fontsize=16)
    
    # 初始化切片索引
    current_slice = [0]  # 使用列表实现闭包内的可变性
    max_slice = img.shape[1] - 1
    
    # 存储图像对象
    image_plots = [[None for _ in range(3)] for _ in range(3)]
    
    # 初始化显示
    for mod in range(3):
        # 原始图像
        image_plots[0][mod] = axs[0, mod].imshow(img[mod, 0], cmap='gray')
        axs[0, mod].set_title(f"{mod_names[mod]} Image")
        
        # ROI掩膜
        image_plots[1][mod] = axs[1, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.5)
        axs[1, mod].set_title(f"{mod_names[mod]} ROI Mask")
        
        # 叠加显示
        image_plots[2][mod] = axs[2, mod].imshow(img[mod, 0], cmap='gray')
        image_plots[2][mod] = axs[2, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.3)
        axs[2, mod].set_title("Overlay")
    
    def update_display():
        """更新所有模态的显示"""
        fig.suptitle(f"Label: {label} | Slice: {current_slice[0]}/{max_slice}", fontsize=16)
        for mod in range(3):
            # 更新原始图像
            image_plots[0][mod].set_data(img[mod, current_slice[0]])
            
            # 更新ROI掩膜
            image_plots[1][mod].set_data(roi[mod, current_slice[0]])
            
            # 更新叠加显示（需要重新创建以避免透明度混合问题）
            axs[2, mod].images[0].set_data(img[mod, current_slice[0]])
            axs[2, mod].images[1].set_data(roi[mod, current_slice[0]])
        
        fig.canvas.draw_idle()
    
    def on_scroll(event):
        """鼠标滚轮事件处理"""
        if event.button == 'up':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.button == 'down':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    # 绑定滚轮事件
    fig.canvas.mpl_connect('scroll_event', on_scroll)
    
    # 添加键盘事件支持
    def on_key(event):
        """键盘事件处理（左右方向键）"""
        if event.key == 'right':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.key == 'left':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    fig.canvas.mpl_connect('key_press_event', on_key)
    
    plt.tight_layout()
    plt.show()

# 使用示例
if __name__ == "__main__":
    dataset = MRIDataset(root_dir='./datasets/pT/train')
    print(len(dataset))
    print(dataset.class_counts) 
    print("有效患者数量:", len(dataset.patient_ids)) 

    loader = DataLoader(dataset, batch_size=1)
    
    img_batch, roi_batch, labels = next(iter(loader))
    visualize_attention(img_batch, roi_batch, labels)

    # 验证数据形状

    print(f"Batch shape: {img_batch.shape}")  # 应该输出 (B, M=3, D=8, H=224, W=224)
    print(f"ROI   shape: {roi_batch.shape}")
